# -*- coding: utf-8 -*-

# @Project : fastapi-tutorial
# @Date    : 20240401-1648
# @Author  : robin

import sys
from dotenv import load_dotenv
load_dotenv("../.env")

#使用openai的官方sdk
import openai
import os

openai.api_key = os.environ.get("OPENAI_API_KEY")

# messages = [
# {"role": "user", "content": "介绍下你自己"}
# ]
#
# res = openai.ChatCompletion.create(
# model="gpt-3.5-turbo",
# messages=messages,
# stream=False,
# )
#
# print(res['choices'][0]['message']['content'])

# Lost in the middle: 长上下文精度问题

# ! pip install sentence-transformers

from langchain.chains import LLMChain,StuffDocumentsChain
from langchain_community.document_transformers import (
    LongContextReorder
)
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain_community.vectorstores import  Chroma

#使用huggingface托管的开源LLM来做嵌入，MiniLM-L6-v2是一个较小的LLM
embedings = HuggingFaceBgeEmbeddings(model_name="all-MiniLM-L6-v2")

text = [
    "篮球是一项伟大的运动。",
    "带我飞往月球是我最喜欢的歌曲之一。",
    "凯尔特人队是我最喜欢的球队。",
    "这是一篇关于波士顿凯尔特人的文件。",
    "我非常喜欢去看电影。",
    "波士顿凯尔特人队以20分的优势赢得了比赛。",
    "这只是一段随机的文字。",
    "《艾尔登之环》是过去15年最好的游戏之一。",
    "L.科内特是凯尔特人队最好的球员之一。",
    "拉里.伯德是一位标志性的NBA球员。"
]

retrieval = Chroma.from_texts(text,embedings).as_retriever(
    search_kwargs={"k": 10}
)

query = "关于我的喜好都知道什么?"


#根据相关性返回文本块
docs = retrieval.get_relevant_documents(query)
print(docs)

#对检索结果进行重新排序，根据论文的方案
#问题相关性越低的内容块放在中间
#问题相关性越高的内容块放在头尾

reordering = LongContextReorder()
reo_docs = reordering.transform_documents(docs)

#头尾共有4个高相关性内容块
print(reo_docs)


from dotenv import load_dotenv
load_dotenv("../.env")
import os
api_key = os.environ.get("OPENAI_API_KEY")
api_base = os.environ.get("OPENAI_API_BASE")

#检测下这种方案的精度效果
from langchain.prompts import PromptTemplate
from langchain.llms import OpenAI

#设置llm
llm = OpenAI(
    api_key=api_key,
    api_base=api_base,
    model="gpt-3.5-turbo-instruct",
    temperature=0
)

document_prompt = PromptTemplate(
    input_variables=["page_content"],template="{page_content}"
)

stuff_prompt_override ="""Given this text extracts:
----------------------------------------
{context}
----------------------------------------
Please answer the following questions:
{query}
"""

prompt = PromptTemplate(
    template=stuff_prompt_override,
    input_variables=["context","query"]
)

llm_chain = LLMChain(
    llm=llm,
    prompt=prompt
)

WorkChain = StuffDocumentsChain(
    llm_chain=llm_chain,
    document_prompt=document_prompt,
    document_variable_name="context"
)

#调用
r = WorkChain.run(
    input_documents=docs,
    query="我最喜欢的球队里最好的球员是谁？"
)
print(r)